/**
 * @file main.cpp
 *
 * Copyright (C) 2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include "data_utils.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
extern void fills_custom_do(uint32_t blockDim, void *stream, uint8_t *src, uint8_t *dst, float negativeSlope);
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void fills_custom(GM_ADDR src, GM_ADDR dst, float negativeSlope);
#endif

int32_t main(int32_t argc, char *argv[])
{
    uint32_t blockDim = 8;
    size_t inputByteSize = 8 * 2048 * sizeof(float);
    size_t outputByteSize = 8 * 2048 * sizeof(float);
    float negativeSlope = 4.8;

#ifdef ASCENDC_CPU_DEBUG
    uint8_t *src = (uint8_t *)AscendC::GmAlloc(inputByteSize);
    uint8_t *dst = (uint8_t *)AscendC::GmAlloc(outputByteSize);

    ReadFile("./input/input_src.bin", inputByteSize, src, inputByteSize);

    AscendC::SetKernelMode(KernelMode::AIV_MODE);
    ICPU_RUN_KF(fills_custom, blockDim, src, dst, negativeSlope); // use this macro for cpu debug

    WriteFile("./output/output_dst.bin", dst, outputByteSize);

    AscendC::GmFree((void *)src);
    AscendC::GmFree((void *)dst);
#else
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    uint8_t *srcHost, *dstHost;
    uint8_t *srcDevice, *dstDevice;

    CHECK_ACL(aclrtMallocHost((void **)(&srcHost), inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&dstHost), outputByteSize));
    CHECK_ACL(aclrtMalloc((void **)&srcDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&dstDevice, outputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

    ReadFile("./input/input_src.bin", inputByteSize, srcHost, inputByteSize);

    CHECK_ACL(aclrtMemcpy(srcDevice, inputByteSize, srcHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    fills_custom_do(blockDim, stream, srcDevice, dstDevice, negativeSlope);
    CHECK_ACL(aclrtSynchronizeStream(stream));

    CHECK_ACL(aclrtMemcpy(dstHost, outputByteSize, dstDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/output_dst.bin", dstHost, outputByteSize);

    CHECK_ACL(aclrtFree(srcDevice));
    CHECK_ACL(aclrtFree(dstDevice));
    CHECK_ACL(aclrtFreeHost(srcHost));
    CHECK_ACL(aclrtFreeHost(dstHost));

    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}
